from torch.nn import Module
from torch.utils.data import Dataset
import logging
import time
import torch


class TrainParameter(object):
    """
    该类用于存储训练器（Trainer)训练时所需的参数信息，包括的基本信息有\n
    epoch: 训练的epoch数\n
    batch_size: 每一次训练使用的数据样本数\n
    num_works: 数据读取时的并发线程数(Windows操作系统中应设置为0)\n
    data_shuffle: 每次训练选取样本时是否随机选取\n
    如果自定义的训练器需要更多参数，可以继承该类以存放更多参数
    """
    def __init__(self, epoch: int, batch_size: int, num_works: int, data_shuffle: bool):
        self.epoch = epoch
        self.batch_size = batch_size
        self.num_works = num_works
        self.data_shuffle = data_shuffle
        self.paras = dict()

    def add_parameter(self, name, value):
        self.paras[name]=value

    def get_parameter(self, name):
        return self.paras.get(name)



class Trainer(object):
    """
    网络训练器。封装了一般网络训练的基本逻辑以及日志记录器(logger)。\n
    note: 自定义训练器时，必须重写的函数有:\n
    train_detail(index:int, item: list) 该函数无需返回任何值。\n
    可选的重写函数有:\n
    epoch_finish() 在每一个epoch训练完成时触发该函数\n
    train_finish() 在整个训练完成时触发该函数\n
    该训练器使用start_train()函数开始训练\n

    Attributes:
        module: 保存当前训练器使用的网络结构(Module)
        dataset: 保存当前网络训练所需使用的数据集

    """
    def __init__(self, module: Module, dataset: Dataset, parameter: TrainParameter):
        self.module = module
        self.dataset = dataset
        self.train_parameter = parameter

        self.logger = None
        self.file_handle = None
        self.console_handle = None
        self.__init_logger()

    def set_module(self, module: Module):
        self.module = module

    def set_dataset(self, dataset: Dataset):
        self.dataset = dataset

    def start_train(self):
        """
        调用该函数以开始网络训练

        :return: None
        """
        self.logger.info('开始训练...')
        try:
            start_time = time.time()
            self.module.train()
            data_loader = torch.utils.data.DataLoader(self.dataset,
                                                      batch_size=self.train_parameter.batch_size,
                                                      shuffle=self.train_parameter.data_shuffle,
                                                      num_workers=self.train_parameter.num_works)
            for train_index in range(0, self.train_parameter.epoch):
                for index, item in enumerate(data_loader):
                    self.train_detail(index, item)
                self.epoch_finish(train_index)
            end_time = time.time()
            self.train_finish()
            self.logger.info("训练已完成，耗时: %d s" % (end_time - start_time))
        except Exception:
            self.logger.error("训练过程失败...", exc_info='True')

    def __init_logger(self):
        self.logger = logging.getLogger("training")
        self.logger.setLevel(logging.INFO)

        self.file_handle = logging.FileHandler("./log.txt", encoding='utf-8')
        self.console_handle = logging.StreamHandler()

        formatter = logging.Formatter('%(asctime)s : %(levelname)s %(message)s')

        self.file_handle.setFormatter(formatter)
        self.console_handle.setFormatter(formatter)

        self.open_console_log_output()
        self.open_file_log_record()

    def close_file_log_record(self):
        """
        关闭日志记录至文件功能（日志文件默认保存至./log.txt)

        :return: None
        """
        self.logger.removeHandler(self.file_handle)

    def open_file_log_record(self):
        """
        开启日志记录至文件功能（日志文件默认保存至./log.txt)

        :return: None
        """
        self.logger.addHandler(self.file_handle)

    def close_console_log_output(self):
        """
        关闭日志输出至控制台功能

        :return: None
        """
        self.logger.removeHandler(self.console_handle)

    def open_console_log_output(self):
        """
        开启日志输出至控制台功能

        :return: None
        """
        self.logger.addHandler(self.console_handle)

    def train_detail(self, index: int, item: list):
        """
        该函数中定义网络训练的细节，如向前传播规则，损失计算，反向传播等。该函数必须被重写

        :param index: 当前训练的batch序号
        :param item: 从数据集中取出的一个batch的数据
        :return: None
        """
        return NotImplementedError()

    def epoch_finish(self, epoch: int):
        """
        当前epoch训练完成时触发的回调函数。

        :param epoch: 当前训练完成的epoch序号
        :return: None
        """
        pass

    def train_finish(self):
        """
        当前整个网络训练完成时触发的回调函数。

        :return: None
        """
        pass
